import sys
import os
import requests
import zipfile
import argparse
import psutil
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QLabel, QProgressBar, QMessageBox
from PyQt5.QtCore import QThread, pyqtSignal, pyqtSlot

class UpdateWorker(QThread):
    update_progress = pyqtSignal(int)
    update_message = pyqtSignal(str)
    update_finished = pyqtSignal()

    def __init__(self, version, check_url, original_program):
        super().__init__()
        self.current_version = version
        self.check_url = check_url
        self.original_program = original_program  # 确保 original_program 被正确设置
        self.download_url = None

    def run(self):
        try:
            # 检查更新
            self.update_message.emit("检查更新中...")
            response = requests.get(self.check_url)
            data = response.json()

            # 获取最新版本号和下载链接
            latest_version = data["release"]["tag"]["name"]
            self.download_url = data["release"]["release"]["attach_files"][0]["cli_download_url"]

            if latest_version and latest_version != self.current_version:
                self.update_message.emit(f"发现新版本: {latest_version}，下载更新中...")
                self.download_update()
            else:
                self.update_message.emit("已是最新版本")
                self.update_finished.emit()

        except Exception as e:
            self.update_message.emit(f"更新失败: {str(e)}")
            self.update_finished.emit()

    def download_update(self):
        response = requests.get(self.download_url, stream=True)
        total_size = int(response.headers.get('content-length', 0))
        downloaded_size = 0

        with open('update.zip', 'wb') as file:
            for data in response.iter_content(chunk_size=1024):
                downloaded_size += len(data)
                file.write(data)
                self.update_progress.emit(int((downloaded_size / total_size) * 100))

        self.update_progress.emit(100)
        self.extract_update()

    def extract_update(self):
        self.update_message.emit("解压更新中...")
        zip_path = "update.zip"
        try:
            if zipfile.is_zipfile(zip_path):
                with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                    zip_ref.extractall("./")
                print(f"解压完成，文件解压到: ./")
                os.remove('update.zip')
                self.update_finished.emit()
            else:
                print("下载的文件不是有效的压缩包")
        except zipfile.BadZipFile as e:
            print(f"压缩包文件损坏: {str(e)}")
        except zipfile.LargeZipFile as e:
            print(f"压缩包过大: {str(e)}")
        except Exception as e:
            print(f"解压时发生错误: {str(e)}")

    def close_original_program(self):
        try:
            for proc in psutil.process_iter(['name']):
                if proc.info['name'] == self.original_program:
                    proc.kill()
                    proc.wait()
                    break
        except Exception as e:
            self.update_message.emit(f"关闭原程序失败: {str(e)}")

    def start_original_program(self):
        try:
            os.startfile(self.original_program)
        except Exception as e:
            self.update_message.emit(f"启动原程序失败: {str(e)}")

class UpdateProgram(QWidget):
    def __init__(self, version, check_url, original_program):
        super().__init__()
        self.init_ui()
        self.worker = UpdateWorker(version, check_url, original_program)

        # 连接信号和槽
        self.worker.update_progress.connect(self.progress.setValue)
        self.worker.update_message.connect(self.label.setText)
        self.worker.update_finished.connect(self.on_update_finished)

        # 启动更新线程
        self.worker.start()

    def init_ui(self):
        self.setWindowTitle('更新程序')
        self.layout = QVBoxLayout()

        self.label = QLabel('初始化...')
        self.layout.addWidget(self.label)

        self.progress = QProgressBar()
        self.layout.addWidget(self.progress)

        self.setLayout(self.layout)
        self.resize(300, 100)
        self.show()

    @pyqtSlot()
    def on_update_finished(self):
        self.worker.close_original_program()
        self.worker.start_original_program()
        QMessageBox.information(self, "更新完成", "更新已成功完成。")
        QApplication.quit()

if __name__ == "__main__":
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="更新程序")
    parser.add_argument("--version", type=str, required=True, help="当前版本号")
    parser.add_argument("--check_url", type=str, required=True, help="检查更新的URL")
    parser.add_argument("--original_program", type=str, required=True, help="原程序的名称或路径")
    args = parser.parse_args()

    app = QApplication(sys.argv)
    update_prog = UpdateProgram(args.version, args.check_url, args.original_program)
    sys.exit(app.exec_())
